/****************************************************************************
 *
 * Copyright (C) 2003-2008 Los Alamos National Security, LLC
 *                         Packet Analytics Corporation
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License Version 2 as
 * published by the Free Software Foundation.  You may not use, modify or
 * distribute this program under any other version of the GNU General
 * Public License.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 *
 ****************************************************************************/
package nfse.stats;

import java.net.*;
import javax.net.*;
import java.io.*;
import java.sql.*;
import java.util.HashMap;
import java.util.Properties;
import java.util.Set;
import java.util.Vector;
import java.util.Iterator;
import java.util.concurrent.*;

import nfse.*;

public class StatsServer extends Thread {
	private int port = -1;

	private Statement statement = null;

	private final int MAX_SIZE = 1000;

	// public final static int INTERVAL = 3600;
	public final static int INTERVAL = 3600;

	public final static int DIRECTION_INCOMING = 1;

	public final static int DIRECTION_OUTGOING = 2;

	public final static int DIRECTION_INTERNAL = 3;

	HashMap<Integer,HashMap<Long,StatsServer.TimeSeries>> timeSeriesBySensor = new HashMap<Integer,HashMap<Long,StatsServer.TimeSeries>>();

	private Connection connection = null;

	public StatsServer(int portNum) {
		this.port = portNum;
	}

	private TimeSeries getTimeSeriesBySensor(int sensorID, Long tsVal) {
		HashMap<Long,StatsServer.TimeSeries> temp = timeSeriesBySensor.get(new Integer(sensorID));
		if (temp == null) {
			temp = new HashMap<Long,StatsServer.TimeSeries>();
			timeSeriesBySensor.put(new Integer(sensorID), temp);
		}
		return (TimeSeries) temp.get(tsVal);
	}

	private void addTimeSeriesBySensor(int sensorID, Long tsVal, TimeSeries ts) {
		HashMap<Long,StatsServer.TimeSeries> temp = timeSeriesBySensor.get(new Integer(sensorID));
		if (temp == null) {
			temp = new HashMap<Long,StatsServer.TimeSeries>();
			timeSeriesBySensor.put(new Integer(sensorID), temp);
		}
		temp.put(tsVal, ts);
	}

	public synchronized void addStatsFile(File f) throws Exception {
		Statement s = connection.createStatement();
		s.execute("insert into stats_files values ('" + f.getAbsolutePath()
				+ "','" + new Timestamp(System.currentTimeMillis()) + "')");
		s.close();
	}

	public synchronized void addStats(int sensorID, StatsEntry stats)
			throws Exception {

		// if (true)
		// return;

		Long tsVal = new Long(((stats.endTS.getTime() / 1000) / INTERVAL)
				* INTERVAL);
		TimeSeries ts = getTimeSeriesBySensor(sensorID, tsVal);

		boolean isNew = false;
		if (ts == null)
			isNew = true;

		if (isNew) {
			ts = new TimeSeries(sensorID, tsVal.longValue());
		} else {

		}

		int direction = stats.direction;

		if (stats.srcIP != -1) {
			ts.addSrcIP(stats.srcIP, direction);
			if (stats.srcBytes != -1) {
				ts.addSrcBytes(stats.srcIP, stats.srcBytes, direction);
			}
			if (stats.srcPackets != -1) {
				ts.addSrcPackets(stats.srcIP, stats.srcPackets, direction);
			}
		}

		if (stats.dstIP != -1) {
			ts.addDstIP(stats.dstIP, direction);
			if (stats.dstBytes != -1) {
				ts.addDstBytes(stats.dstIP, stats.dstBytes, direction);
			}
			if (stats.dstPackets != -1) {
				ts.addDstPackets(stats.dstIP, stats.dstPackets, direction);
			}
		}

		if (stats.dstPort != -1) {
			ts.addDstPort(stats.dstPort, direction);
		}

		if (stats.protocol != -1) {
			ts.addProtocol(stats.protocol, direction);
		}

		if (isNew) {
			addTimeSeriesBySensor(sensorID, tsVal, ts);

		} else {
			if (ts.sizeExceeded()) {

				this.update(statement);
			}
		}
	}

	public synchronized TimeSeries[] getTimeSeries(int sensorID) {
		HashMap<Long,StatsServer.TimeSeries> timeSeries = timeSeriesBySensor.get(new Integer(
				sensorID));
		TimeSeries[] temp = new TimeSeries[timeSeries.size()];
		Iterator<Long> it = timeSeries.keySet().iterator();
		int idx = 0;
		while (it.hasNext()) {
			temp[idx] = (TimeSeries) timeSeries.get((Long) it.next());
			idx++;
		}
		return temp;
	}

	public static synchronized void addRunStats(RunStats stats) {
		try {
			Properties props = new Properties();
			props.setProperty("user", NetFSE.STATS_DB_USER);
			props.setProperty("password", NetFSE.STATS_DB_PASSWORD);
			props.setProperty("client_encoding", "UTF8");

			Class.forName("org.postgresql.Driver");
			String url = "jdbc:postgresql://" + NetFSE.STATS_DB_HOST + "/"
					+ NetFSE.STATS_DB_DATABASE;

			Connection connection = DriverManager.getConnection(url, props);
			Statement statement = connection.createStatement();
			String sql = "select * from Sensor_Stats where Sensor_ID="
					+ stats.sensorID;
			ResultSet rs = statement.executeQuery(sql);
			if (rs.next()) {
				Timestamp minTS = rs.getTimestamp("Min_TS");
				Timestamp maxTS = rs.getTimestamp("Max_TS");
				long count = rs.getLong("Record_Count");
				if (stats.minTS.before(minTS))
					minTS = stats.minTS;
				if (stats.maxTS.before(maxTS))
					maxTS = stats.maxTS;
				count += stats.count;
				sql = "update Sensor_Stats set Min_TS='" + minTS
						+ "', Max_TS='" + maxTS + "', Record_Count=" + count
						+ " where Sensor_ID=" + stats.sensorID
						+ " and Type_ID=" + stats.typeID;
			} else {
				sql = "insert into Sensor_Stats values (" + stats.typeID + ","
						+ stats.sensorID + ",'" + stats.minTS + "','"
						+ stats.maxTS + "'," + stats.count + ")";
			}
			rs.close();
			statement.execute(sql);
			statement.close();
			connection.close();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public synchronized void update(Statement statement) throws Exception {

		System.out.println("Started update process at "
				+ new Timestamp(System.currentTimeMillis()));

		File f = File.createTempFile("nfse_stats_", ".sql");

		Iterator<Integer> it2 = timeSeriesBySensor.keySet().iterator();
		while (it2.hasNext()) {
			HashMap<Long,StatsServer.TimeSeries> timeSeries = this.timeSeriesBySensor.get(it2
					.next());

			TimeSeries[] temp = new TimeSeries[timeSeries.size()];
			Iterator<Long> it = timeSeries.keySet().iterator();
			int idx = 0;
			while (it.hasNext()) {
				temp[idx] = (TimeSeries) timeSeries.get((Long) it.next());
				temp[idx].update(statement, f);
				idx++;
			}
			for (int i = 0; i < temp.length; i++) {
				timeSeries.remove(temp[i].ts);
			}
		}

	}

	public void run() {
		try {
			System.out.println("Stats Server: Listening on port " + port + ".");
			ServerSocketFactory ssf = null;
			ssf = ServerSocketFactory.getDefault();
			ServerSocket ss = ssf.createServerSocket(port);
			// address = ss.getInetAddress();

			Properties props = new Properties();
			props.setProperty("user", NetFSE.STATS_DB_USER);
			props.setProperty("password", NetFSE.STATS_DB_PASSWORD);
			props.setProperty("client_encoding", "UTF8");

			Class.forName("org.postgresql.Driver");
			String url = "jdbc:postgresql://" + NetFSE.STATS_DB_HOST + "/"
					+ NetFSE.STATS_DB_DATABASE;

			connection = DriverManager.getConnection(url, props);

			this.statement = connection.createStatement();

			RefreshThread rt = new RefreshThread(this, connection
					.createStatement());
			rt.start();

			StatsSocketMonitor dsm = new StatsSocketMonitor(300);
			dsm.start();

			while (true) {
				StatsSocket sock = new StatsSocket(ss.accept(), this);
				dsm.addStatsSocket(sock);
				sock.start();
				System.out.println("Stats Server: Opened new stats socket.");
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public class TimeSeries {

		private final Semaphore available = new Semaphore(1, true);

		HashMap[] dstIpMaps = new HashMap[3];

		HashMap[] srcIpMaps = new HashMap[3];

		HashMap[] dstPortMaps = new HashMap[3];

		HashMap[] protocolMaps = new HashMap[3];

		HashMap[] dstBytesMaps = new HashMap[3];

		HashMap[] srcBytesMaps = new HashMap[3];

		HashMap[] dstPacketsMaps = new HashMap[3];

		HashMap[] srcPacketsMaps = new HashMap[3];

		public long ts = 0;

		// public short direction = 1;

		public int sensorID = -1;

		public TimeSeries(int sensorID, long ts) {
			this.ts = ts;
			this.sensorID = sensorID;
			for (int i = 0; i < dstIpMaps.length; i++) {
				dstIpMaps[i] = new HashMap(100);
				srcIpMaps[i] = new HashMap(100);
				protocolMaps[i] = new HashMap(100);
				dstPortMaps[i] = new HashMap(100);
				dstBytesMaps[i] = new HashMap(100);
				srcBytesMaps[i] = new HashMap(100);
				dstPacketsMaps[i] = new HashMap(100);
				srcPacketsMaps[i] = new HashMap(100);
			}
		}

		public synchronized boolean sizeExceeded() {
			for (int i = 0; i < dstIpMaps.length; i++) {
				if (dstIpMaps[i].size() > MAX_SIZE)
					return true;
				if (srcIpMaps[i].size() > MAX_SIZE)
					return true;
				if (dstPortMaps[i].size() > MAX_SIZE)
					return true;
				if (protocolMaps[i].size() > MAX_SIZE)
					return true;
				if (dstBytesMaps[i].size() > MAX_SIZE)
					return true;
				if (dstPacketsMaps[i].size() > MAX_SIZE)
					return true;
				if (srcBytesMaps[i].size() > MAX_SIZE)
					return true;
				if (srcPacketsMaps[i].size() > MAX_SIZE)
					return true;
			}
			return false;
		}

		public synchronized void addSrcIP(long srcIP, int direction) {
			try {
				available.acquire();
				// long srcIPa = srcIP & 0xFF000000;
				// long srcIPb = srcIP & 0xFFFF0000;
				long srcIPc = srcIP & 0xFFFFFF00;

				if (!NetFSE.isHomeNet(srcIP))
					srcIP = srcIPc;

				String srcAddr = NetFSEUtil.convertLongToIP(srcIP);

				HashMap map = srcIpMaps[direction - 1];

				Long val = (Long) map.get(srcAddr);
				if (val != null) {
					map.put(srcAddr, new Long(val.longValue() + 1));
				} else {
					map.put(srcAddr, new Long(1));
				}
				available.release();
			} catch (Exception e) {
				e.printStackTrace();
			}
		}

		public synchronized void addDstIP(long dstIP, int direction) {
			try {
				available.acquire();
				// long dstIPa = dstIP & 0xFF000000;
				// long dstIPb = dstIP & 0xFFFF0000;
				long dstIPc = dstIP & 0xFFFFFF00;

				if (!NetFSE.isHomeNet(dstIP))
					dstIP = dstIPc;

				String dstAddr = NetFSEUtil.convertLongToIP(dstIP);

				HashMap map = dstIpMaps[direction - 1];

				Long val = (Long) map.get(dstAddr);
				if (val != null) {
					map.put(dstAddr, new Long(val.longValue() + 1));
				} else {
					map.put(dstAddr, new Long(1));
				}
				available.release();
			} catch (Exception e) {
				e.printStackTrace();
			}
		}

		public synchronized void addDstPort(int dstPort, int direction) {
			try {
				// if (dstPort > 1024)
				// return;
				available.acquire();

				HashMap map = dstPortMaps[direction - 1];

				Long val = (Long) map.get(new Integer(dstPort));
				if (val != null) {
					map
							.put(new Integer(dstPort), new Long(
									val.longValue() + 1));
				} else {
					map.put(new Integer(dstPort), new Long(1));
				}
				available.release();
			} catch (Exception e) {
				e.printStackTrace();
			}
		}

		public synchronized void addProtocol(short protocol, int direction) {
			try {
				available.acquire();

				HashMap map = protocolMaps[direction - 1];

				Long val = (Long) map.get(new Short(protocol));
				if (val != null) {
					map.put(new Short(protocol), new Long(val.longValue() + 1));
				} else {
					map.put(new Short(protocol), new Long(1));
				}
				available.release();
			} catch (Exception e) {
				e.printStackTrace();
			}
		}

		public synchronized void addDstPackets(long dstIP, long dstPackets,
				int direction) {
			try {
				available.acquire();

				long dstIPc = dstIP & 0xFFFFFF00;
				String dstAddr = NetFSEUtil.convertLongToIP(dstIP);

				if (!NetFSE.isHomeNet(dstIP))
					dstIP = dstIPc;

				HashMap map = dstBytesMaps[direction - 1];

				Long val = (Long) map.get(dstAddr);
				if (val != null) {
					map.put(dstAddr, new Long(val.longValue() + dstPackets));
				} else {
					map.put(dstAddr, new Long(dstPackets));
				}
				available.release();
			} catch (Exception e) {
				e.printStackTrace();
			}
		}

		public synchronized void addSrcPackets(long srcIP, long srcPackets,
				int direction) {
			try {
				available.acquire();

				long srcIPc = srcIP & 0xFFFFFF00;
				String srcAddr = NetFSEUtil.convertLongToIP(srcIP);

				if (!NetFSE.isHomeNet(srcIP))
					srcIP = srcIPc;

				HashMap map = srcPacketsMaps[direction - 1];

				Long val = (Long) map.get(srcAddr);
				if (val != null) {
					map.put(srcAddr, new Long(val.longValue() + srcPackets));
				} else {
					map.put(srcAddr, new Long(srcPackets));
				}
				available.release();
			} catch (Exception e) {
				e.printStackTrace();
			}
		}

		public synchronized void addDstBytes(long dstIP, long dstBytes,
				int direction) {
			try {
				available.acquire();

				long dstIPc = dstIP & 0xFFFFFF00;
				String dstAddr = NetFSEUtil.convertLongToIP(dstIP);

				if (!NetFSE.isHomeNet(dstIP))
					dstIP = dstIPc;

				HashMap map = dstBytesMaps[direction - 1];

				Long val = (Long) map.get(dstAddr);
				if (val != null) {
					map.put(dstAddr, new Long(val.longValue() + dstBytes));
				} else {
					map.put(dstAddr, new Long(dstBytes));
				}
				available.release();
			} catch (Exception e) {
				e.printStackTrace();
			}
		}

		public synchronized void addSrcBytes(long srcIP, long srcBytes,
				int direction) {
			try {
				available.acquire();

				long srcIPc = srcIP & 0xFFFFFF00;
				String srcAddr = NetFSEUtil.convertLongToIP(srcIP);

				if (!NetFSE.isHomeNet(srcIP))
					srcIP = srcIPc;

				HashMap map = srcBytesMaps[direction - 1];

				Long val = (Long) map.get(srcBytes);
				if (val != null) {
					map.put(srcAddr, new Long(val.longValue() + srcBytes));
				} else {
					map.put(srcAddr, new Long(srcBytes));
				}
				available.release();
			} catch (Exception e) {
				e.printStackTrace();
			}
		}

		public int[] getSizes() {
			int[] temp = new int[0];

			return temp;
		}

		public Timestamp getTimestamp() {
			return new Timestamp(ts * 1000);
		}

		public synchronized void update(Statement statement, File f) {
			try {
				available.acquire();

				HashMap[][] maps = new HashMap[8][];
				String[] tables = new String[maps.length];
				String[] fields = new String[maps.length];
				String[] statNames = new String[maps.length];
				boolean[] isString = new boolean[maps.length];

				maps[0] = dstIpMaps;
				maps[1] = srcIpMaps;
				maps[2] = dstPortMaps;
				maps[3] = protocolMaps;
				maps[4] = dstBytesMaps;
				maps[5] = srcBytesMaps;
				maps[6] = dstPacketsMaps;
				maps[7] = srcPacketsMaps;

				tables[0] = "Dst_IP_Stats";
				tables[1] = "Src_IP_Stats";
				tables[2] = "Dst_Port_Stats";
				tables[3] = "Protocol_Stats";
				tables[4] = "Dst_Bytes_Stats";
				tables[5] = "Src_Bytes_Stats";
				tables[6] = "Dst_Packets_Stats";
				tables[7] = "Src_Packets_Stats";

				fields[0] = "IP";
				fields[1] = "IP";
				fields[2] = "Port";
				fields[3] = "Protocol";
				fields[4] = "IP";
				fields[5] = "IP";
				fields[6] = "IP";
				fields[7] = "IP";

				statNames[0] = "Num_Hits";
				statNames[1] = "Num_Hits";
				statNames[2] = "Num_Hits";
				statNames[3] = "Num_Hits";
				statNames[4] = "Num_Bytes";
				statNames[5] = "Num_Bytes";
				statNames[6] = "Num_Packets";
				statNames[7] = "Num_Packets";

				isString[0] = true;
				isString[1] = true;
				isString[2] = false;
				isString[3] = false;
				isString[4] = true;
				isString[5] = true;
				isString[6] = true;
				isString[7] = true;

				for (int x = 0; x < maps.length; x++) {

					for (int y = 0; y < maps[x].length; y++) {

						int direction = y + 1;
						Vector<String> values = new Vector<String>();

						String sql = "select * from " + tables[x]
								+ " where Direction=" + direction
								+ " and Sensor_ID=" + sensorID
								+ " and Time_Interval=" + this.ts;
						// System.out.println(sql);
						ResultSet rs = statement.executeQuery(sql);
						while (rs.next()) {
							Object value = null;
							Long hits = rs.getLong(statNames[x]);

							if (x == 3) {
								Short protocol = rs.getShort(fields[x]);
								value = protocol;
							} else if (!isString[x]) {
								Integer port = rs.getInt(fields[x]);
								value = port;
								// if (port > 1024)
								// value = null;
							} else {
								String ip = rs.getString(fields[x]);
								value = ip;
							}

							Long hits2 = (Long) maps[x][y].get(value);
							if (hits2 != null) {
								maps[x][y].remove(value);
							} else {
								hits2 = new Long(0);
							}
							hits = hits + hits2;
							if (value != null) {
								String insert = "";
								if (isString[x])
									insert = sensorID + "," + ts + ",'" + value
											+ "'," + hits + "," + direction;
								else
									insert = sensorID + "," + ts + "," + value
											+ "," + hits + "," + direction;
								values.addElement(insert);
							}
						}
						rs.close();

						Set keys = maps[x][y].keySet();
						Object[] keyArray = keys.toArray();
						for (int i = 0; i < keyArray.length; i++) {
							Long hits = (Long) maps[x][y].get(keyArray[i]);
							String insert = "";
							if (isString[x])
								insert = sensorID + "," + ts + ",'"
										+ keyArray[i] + "'," + hits + ","
										+ direction;
							else
								insert = sensorID + "," + ts + ","
										+ keyArray[i] + "," + hits + ","
										+ direction;
							values.addElement(insert);
							maps[x][y].remove(keyArray[i]);
						}

						if (values.size() > 0) {

							// PrintWriter out = new PrintWriter(new
							// FileWriter(f,
							// true));

							statement.execute("delete from " + tables[x]
									+ " where Direction = " + direction
									+ " and Sensor_ID=" + sensorID
									+ " and Time_Interval=" + ts);

							for (int i = 0; i < values.size(); i++) {
								String insertSQL = "insert  into " + tables[x]
										+ " values (" + values.elementAt(i)
										+ ")";
								statement.execute(insertSQL);
							}

							// out.flush();
							// out.close();

						}
					}
				}

				available.release();
			} catch (Exception e) {
				e.printStackTrace();
			}
		}
	}

	public class RefreshThread extends Thread {

		private final int INTERVAL = 1;

		private StatsServer server = null;

		private Statement statement = null;

		public RefreshThread(StatsServer server, Statement statement) {
			this.server = server;
			this.statement = statement;
		}

		public void run() {
			try {
				int interval = 1000 * INTERVAL * 60;

				while (true) {
					long startTS = System.currentTimeMillis();

					String sql = "select * from stats_files order by ts asc limit 1";
					ResultSet rs = statement.executeQuery(sql);
					if (rs.next()) {
						String path = rs.getString("path");
						rs.close();
						System.out.println("Processing stats file '" + path
								+ "' ");
						File f = new File(path);
						if (!f.exists()) {
							statement
									.execute("delete from stats_files where path='"
											+ path + "'");
						} else {
							LineNumberReader in = new LineNumberReader(
									new FileReader(f));
							String line = null;
							while ((line = in.readLine()) != null) {
								try {
									StatsEntry stats = StatsEntry
											.generate(line);
									server.addStats(stats.sensorID, stats);
								} catch (Exception e) {
									e.printStackTrace();
									System.out.println("BAD LINE: '" + line
											+ "'");
								}
							}
							in.close();
							server.update(statement);
							// if (!oops) {
							statement
									.execute("delete from stats_files where path='"
											+ path + "'");
							f.delete();
							// }
							long endTS = System.currentTimeMillis();
							System.out.println("Stats file '"
									+ f.getAbsolutePath() + "' processed in "
									+ (endTS - startTS) / 1000 + " seconds");
						}
					} else {
						System.out.println("No stats to process; sleeping");
						Thread.sleep(interval);

					}

				}

			} catch (Exception E) {
				E.printStackTrace();
			}
		}

	}

}